#
# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# openUBMC is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#         http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.
#

"""
嵌套函数控制流生成器

生成包含 await 的控制流语句（for、while、if等）
"""

from typing import Optional
from mcpy.ir.statements.control import IRIf
from mcpy.ir.expressions.basic import IRName
from mcpy.ir.statements.assignment import IRAssign
from mcpy.ir.statements.control import (
    IRReturn,
    IRIf,
    IRForRange,
    IRForEach,
    IRWhile,
    IRBreak,
    IRContinue,
)
from mcpy.ir.expressions.await_expr import IRAwait
from mcpy.codegen.core.context import CodeGenContext
from mcpy.codegen.core.generator import get_global_generator
from .utils import fix_global_async_struct_call
from .type_inference import get_await_result_type


class ControlFlowGenerator:
    """控制流生成器"""

    def __init__(self):
        pass

    def generate_for_with_await(
        self,
        for_stmt,
        state_vars: dict,
        context: CodeGenContext,
        base_indent: str,
        func_scope_local_vars: Optional[set] = None,
    ) -> str:
        """生成包含 await 的 for 循环代码

        Args:
            for_stmt: for 循环节点（IRForRange 或 IRForEach）
            state_vars: 成员变量字典
            context: 代码生成上下文
            base_indent: 基础缩进
            func_scope_local_vars: 函数头声明的局部变量集合

        Returns:
            生成的代码
        """

        if func_scope_local_vars is None:
            func_scope_local_vars = set()

        generator = get_global_generator()
        lines = []

        if isinstance(for_stmt, IRForRange):
            # for i in range(...) 转换为 C++ for 循环
            var_name = for_stmt.var.id if isinstance(for_stmt.var, IRName) else "i"

            # 生成 range 参数
            start_code = (
                generator.generate_expression(for_stmt.start, context)
                if for_stmt.start
                else "0"
            )
            end_code = generator.generate_expression(for_stmt.end, context)
            step_code = (
                generator.generate_expression(for_stmt.step, context)
                if for_stmt.step
                else "1"
            )

            # 检查循环变量是否已经是成员变量
            is_member_var = var_name in state_vars

            if is_member_var:
                # 循环变量是成员变量，需要先初始化再循环
                lines.append(f"{base_indent}{var_name} = {start_code};")
                if for_stmt.step:
                    for_header = (
                        f"for (; {var_name} < {end_code}; {var_name} += {step_code})"
                    )
                else:
                    for_header = f"for (; {var_name} < {end_code}; ++{var_name})"
                lines.append(f"{base_indent}{for_header} {{")
            else:
                # 循环变量是局部变量，直接在 for 中声明
                if for_stmt.step:
                    for_header = f"for (int {var_name} = {start_code}; {var_name} < {end_code}; {var_name} += {step_code})"
                else:
                    for_header = f"for (int {var_name} = {start_code}; {var_name} < {end_code}; ++{var_name})"
                lines.append(f"{base_indent}{for_header} {{")

            # 递归生成循环体
            if for_stmt.body:
                body_code = self.generate_statements_with_await(
                    for_stmt.body.statements,
                    state_vars,
                    context,
                    base_indent + "    ",
                    func_scope_local_vars,
                )
                if body_code:
                    lines.append(body_code)

            lines.append(f"{base_indent}}}")

        elif isinstance(for_stmt, IRForEach):
            # for item in iterable（暂时简单处理）
            var_name = for_stmt.var.id if isinstance(for_stmt.var, IRName) else "item"
            iter_code = generator.generate_expression(for_stmt.iter_expr, context)

            lines.append(f"{base_indent}for (auto {var_name} : {iter_code}) {{")

            # 递归生成循环体
            if for_stmt.body:
                body_code = self.generate_statements_with_await(
                    for_stmt.body.statements,
                    state_vars,
                    context,
                    base_indent + "    ",
                    func_scope_local_vars,
                )
                if body_code:
                    lines.append(body_code)

            lines.append(f"{base_indent}}}")

        return "\n".join(lines)

    def generate_while_with_await(
        self,
        while_stmt,
        state_vars: dict,
        context: CodeGenContext,
        base_indent: str,
        func_scope_local_vars: Optional[set] = None,
        parent_loop_label: Optional[str] = None,
        parent_continue_label: Optional[str] = None,
    ) -> str:
        """生成包含 await 的 while 循环代码

        Args:
            while_stmt: while 循环节点（IRWhile）
            state_vars: 成员变量字典
            context: 代码生成上下文
            base_indent: 基础缩进
            func_scope_local_vars: 函数头声明的局部变量集合
            parent_loop_label: 父循环的退出标签（用于处理嵌套循环中的 break）
            parent_continue_label: 父循环的继续标签（用于处理嵌套循环中的 continue）

        Returns:
            生成的代码
        """
        if func_scope_local_vars is None:
            func_scope_local_vars = set()

        generator = get_global_generator()
        lines = []

        # 生成 while 条件
        condition_code = generator.generate_expression(while_stmt.test, context)
        lines.append(f"{base_indent}while ({condition_code}) {{")

        # 递归生成循环体
        # 注意：现在使用 co_break 和 co_continue，不再需要传递标签
        body_code = self.generate_statements_with_await(
            while_stmt.body.statements,
            state_vars,
            context,
            base_indent + "    ",
            func_scope_local_vars,
            None,  # current_loop_label 不再需要
            None,  # current_continue_label 不再需要
        )
        if body_code:
            lines.append(body_code)

        lines.append(f"{base_indent}}}")

        return "\n".join(lines)

    def generate_if_with_await(
        self,
        if_stmt: IRIf,
        state_vars: dict,
        context: CodeGenContext,
        base_indent: str,
        func_scope_local_vars: Optional[set] = None,
        current_loop_label: Optional[str] = None,
        current_continue_label: Optional[str] = None,
    ) -> str:
        """生成包含 await 的 if 语句

        递归处理 if 语句内部的 await，保持控制流结构

        Args:
            if_stmt: if 语句节点
            state_vars: 成员变量字典
            context: 代码生成上下文
            base_indent: 基础缩进
            func_scope_local_vars: 函数头声明的局部变量集合
            current_loop_label: 当前循环的退出标签（用于 break 转换）
            current_continue_label: 当前循环的继续标签（用于 continue 转换）

        Returns:
            生成的代码
        """
        if func_scope_local_vars is None:
            func_scope_local_vars = set()

        generator = get_global_generator()
        lines = []

        # 生成条件
        condition_code = generator.generate_expression(if_stmt.test, context)
        lines.append(f"{base_indent}if ({condition_code}) {{")

        # 处理 then 分支
        if if_stmt.then_body:
            then_code = self.generate_statements_with_await(
                if_stmt.then_body.statements,
                state_vars,
                context,
                base_indent + "    ",
                func_scope_local_vars,
                current_loop_label,
                current_continue_label,
            )
            lines.append(then_code)

        # 处理 else 分支
        if if_stmt.else_body:
            lines.append(f"{base_indent}}} else {{")
            else_code = self.generate_statements_with_await(
                if_stmt.else_body.statements,
                state_vars,
                context,
                base_indent + "    ",
                func_scope_local_vars,
                current_loop_label,
                current_continue_label,
            )
            lines.append(else_code)

        lines.append(f"{base_indent}}}")
        return "\n".join(lines)

    def generate_statements_with_await(
        self,
        statements: list,
        state_vars: dict,
        context: CodeGenContext,
        indent: str,
        func_scope_local_vars: Optional[set] = None,
        current_loop_label: Optional[str] = None,
        current_continue_label: Optional[str] = None,
    ) -> str:
        """递归生成包含 await 的语句列表

        这是旧版实现，保留用于兼容性。新代码应该使用 async_code_generator。

        Args:
            statements: 语句列表
            state_vars: 成员变量字典
            context: 代码生成上下文
            indent: 缩进
            func_scope_local_vars: 函数头声明的局部变量集合
            current_loop_label: 当前循环的退出标签（用于 break 转换）
            current_continue_label: 当前循环的继续标签（用于 continue 转换）

        Returns:
            生成的代码
        """
        if func_scope_local_vars is None:
            func_scope_local_vars = set()

        generator = get_global_generator()
        lines = []

        for stmt in statements:
            if isinstance(stmt, IRAssign) and isinstance(stmt.value, IRAwait):
                # await 赋值语句
                target = stmt.get_target_expr()

                target_name = target.id if isinstance(target, IRName) else "_"
                await_expr = stmt.value.value
                await_expr_code = generator.generate_expression(await_expr, context)

                # 修正全局异步函数 struct 的调用语法
                # 需要传递当前函数节点以便查找符号
                current_func = context.get_extension_param("current_node")
                if current_func:
                    await_expr_code = fix_global_async_struct_call(
                        await_expr, await_expr_code, current_func
                    )

                # 生成 co_yield_return await(...)
                lines.append(
                    f"{indent}co_yield_return await({await_expr_code}, __ctx);"
                )

                # 生成结果获取（只有当目标变量不是 "_" 时才生成）
                if target_name != "_":
                    if target_name in state_vars:
                        await_type = state_vars[target_name]
                    else:
                        await_type = get_await_result_type(target)
                        if not await_type:
                            await_type = "mc::variant"

                    if target_name in state_vars:
                        lines.append(
                            f"{indent}{target_name} = __resume_value.get<{await_type}>();"
                        )
                    else:
                        lines.append(
                            f"{indent}{await_type} {target_name} = __resume_value.get<{await_type}>();"
                        )

            elif isinstance(stmt, IRReturn):
                # return 语句
                return_expr = stmt.expr
                if return_expr:
                    return_code = generator.generate_expression(return_expr, context)
                    lines.append(f"{indent}return set_result({return_code});")
                else:
                    lines.append(f"{indent}return set_result();")

            elif isinstance(stmt, IRBreak):
                # break 语句：在包含 await 的异步循环中使用 co_break
                # _generate_statements_with_await 用于生成包含 await 的语句，
                # 因此这里的 break 应该使用 co_break
                lines.append(f"{indent}co_break;")

            elif isinstance(stmt, IRContinue):
                # continue 语句：在包含 await 的异步循环中使用 co_continue
                # _generate_statements_with_await 用于生成包含 await 的语句，
                # 因此这里的 continue 应该使用 co_continue
                lines.append(f"{indent}co_continue;")

            elif isinstance(stmt, IRIf):
                # 嵌套的 if 语句
                if_code = self.generate_if_with_await(
                    stmt,
                    state_vars,
                    context,
                    indent,
                    func_scope_local_vars,
                    current_loop_label,
                    current_continue_label,
                )
                lines.append(if_code)

            elif isinstance(stmt, (IRForRange, IRForEach)):
                # for 循环语句
                for_code = self.generate_for_with_await(
                    stmt, state_vars, context, indent, func_scope_local_vars
                )
                lines.append(for_code)

            elif isinstance(stmt, IRWhile):
                # while 循环语句
                while_code = self.generate_while_with_await(
                    stmt,
                    state_vars,
                    context,
                    indent,
                    func_scope_local_vars,
                    current_loop_label,
                    current_continue_label,
                )
                lines.append(while_code)

            elif isinstance(stmt, IRAssign):
                # 检查是否是元组解包赋值
                if hasattr(stmt, "is_tuple_unpacking") and stmt.is_tuple_unpacking:
                    # 元组解包：使用正常的语句生成流程
                    stmt_code = generator.generate_statement(stmt, context)
                    if stmt_code and not stmt_code.isspace():
                        lines.append(f"{indent}{stmt_code}")
                else:
                    # 单目标赋值
                    target = stmt.get_target_expr()
                    if target and isinstance(target, IRName):
                        var_name = target.id
                        if var_name in state_vars or var_name in func_scope_local_vars:
                            # 成员变量或函数头局部变量：只生成赋值，不带类型
                            value_code = generator.generate_expression(
                                stmt.value, context
                            )
                            lines.append(f"{indent}{var_name} = {value_code};")
                        else:
                            # 新的局部变量：按原样生成（带类型）
                            stmt_code = generator.generate_statement(stmt, context)
                            if stmt_code and not stmt_code.isspace():
                                lines.append(f"{indent}{stmt_code}")
                    else:
                        stmt_code = generator.generate_statement(stmt, context)
                        if stmt_code and not stmt_code.isspace():
                            lines.append(f"{indent}{stmt_code}")

            else:
                # 其他语句
                stmt_code = generator.generate_statement(stmt, context)
                if stmt_code and not stmt_code.isspace():
                    lines.append(f"{indent}{stmt_code}")

        return "\n".join(lines)
